package org.geekang.bigdata.hbase;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.hbase.HBaseConfiguration;
import org.apache.hadoop.hbase.TableName;
import org.apache.hadoop.hbase.client.*;
import org.apache.hadoop.hbase.util.Bytes;

import java.io.IOException;
import java.util.Arrays;
import java.util.List;
import java.util.stream.Collectors;

public class HBaseRepository {
    private static HBaseAdmin admin = null;
    private static Connection connection = null;
    private final String ZK_CONNECT_KEY = "hbase.zookeeper.quorum";
    private final String ZK_CONNECT_VAL = "47.101.204.23,47.101.216.12,47.101.206.249";

    /**
     * 初始化连接和Admin对象
     *
     * @throws IOException IO异常
     */
    public void init() throws IOException {
        Configuration conf = HBaseConfiguration.create();
        conf.set(ZK_CONNECT_KEY, ZK_CONNECT_VAL);

        connection = ConnectionFactory.createConnection(conf);
        admin = (HBaseAdmin) connection.getAdmin();
    }

    /**
     * 创建表
     *
     * @param tableName 表名
     * @param colFamily 列族
     * @throws IOException IO异常
     */
    public void createTable(String tableName, String[] colFamily) throws IOException {
        TableName table = TableName.valueOf(tableName);
        if (admin.tableExists(table)) {
            System.out.println(tableName + "表已存在");

            return;
        }

        List<ColumnFamilyDescriptor> descriptorList = Arrays.asList(colFamily).stream().map(
                        col -> ColumnFamilyDescriptorBuilder.newBuilder(Bytes.toBytes(col)).build())
                .collect(Collectors.toList());

        TableDescriptor tableDescriptor = TableDescriptorBuilder.newBuilder(table)
                .setColumnFamilies(descriptorList)
                .build();

        admin.createTable(tableDescriptor);
    }

    /**
     * 删除表
     *
     * @param tableName 表名
     * @throws IOException IO异常
     */
    public void deleteTable(String tableName) throws IOException {
        TableName t = TableName.valueOf(tableName);
        admin.disableTable(t);
        admin.deleteTable(t);
    }

    /**
     * 插入或更新记录
     *
     * @param tableName 表名
     * @param rowKey    键
     * @param colFamily 列族
     * @param column    字段
     * @param value     值
     * @throws IOException IO异常
     */
    public void insertOrUpdate(String tableName, String rowKey, String colFamily, String column, String value) throws IOException {
        this.insertOrUpdate(tableName, rowKey, colFamily, new String[]{column}, new String[]{value});
    }

    /**
     * 插入或更新记录
     *
     * @param tableName 表名
     * @param rowKey    键
     * @param colFamily 列族
     * @param columns   字段
     * @param values    值
     * @throws IOException IO异常
     */
    public void insertOrUpdate(String tableName, String rowKey, String colFamily, String[] columns, String[] values) throws IOException {
        Table table = connection.getTable(TableName.valueOf(tableName));
        Put put = new Put(Bytes.toBytes(rowKey));
        for (int i = 0; i < columns.length; i++) {
            put.addColumn(Bytes.toBytes(colFamily), Bytes.toBytes(columns[i]), Bytes.toBytes(values[i]));
            table.put(put);
        }
    }

    /**
     * 打印表数据
     *
     * @param tableName 表名
     * @throws IOException IO异常
     */
    public void printTable(String tableName) throws IOException {
        Table t = connection.getTable(TableName.valueOf(tableName));
        Scan scan = new Scan();
        ResultScanner scanner = t.getScanner(scan);

        for (Result result : scanner) {
            System.out.println(result);
        }
    }

    /**
     * 关闭连接
     */
    public void close() {
        try {
            if (admin != null) {
                admin.close();
            }
            if (connection != null) {
                connection.close();
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
}
